{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### trainer的官方api https://huggingface.co/docs/transformers/main_classes/trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 导入相关包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments, BertTokenizer, BertForSequenceClassification\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrain_model_dir = \"/data/models/huggingface/rbt3\"\n",
    "save_dir = '/data/logs/data_parallel_origin'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载数据集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d2fa24799d143acbdd4b028295ae918",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/7766 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['label', 'review'],\n",
       "    num_rows: 7765\n",
       "})"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = load_dataset(\"csv\", data_files=\"./ChnSentiCorp_htl_all.csv\", split=\"train\")\n",
    "dataset = dataset.filter(lambda x: x[\"review\"] is not None)\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 划分数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['label', 'review'],\n",
       "        num_rows: 6988\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['label', 'review'],\n",
       "        num_rows: 777\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "datasets = dataset.train_test_split(test_size=0.1)\n",
    "datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建dataloader\n",
    "import torch\n",
    "tokenizer = BertTokenizer.from_pretrained(pretrain_model_dir,device=\"cuda:0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "803aac2deb574567a9c51fe887776882",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/6988 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "439e45f8556140c1ac298b1d6d08a957",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/777 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 6988\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['input_ids', 'token_type_ids', 'attention_mask', 'labels'],\n",
       "        num_rows: 777\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def process_function(examples):\n",
    "    tokenized_examples = tokenizer(examples[\"review\"], max_length=128, truncation=True)\n",
    "    tokenized_examples[\"labels\"] = examples[\"label\"]\n",
    "    return tokenized_examples\n",
    "\n",
    "tokenized_datasets = datasets.map(process_function, batched=True, remove_columns=datasets[\"train\"].column_names)\n",
    "tokenized_datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# dataloader 处理不需要了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader \n",
    "# 这里才正式的开始进行padding操作\n",
    "from transformers import DataCollatorWithPadding\n",
    "\n",
    "trainset,valset  =  tokenized_datasets['train'],tokenized_datasets['test']\n",
    "train_loader = DataLoader(trainset, batch_size=32, shuffle=True, collate_fn=DataCollatorWithPadding(tokenizer=tokenizer))\n",
    "val_loader = DataLoader(valset, batch_size=64, shuffle=False, collate_fn=DataCollatorWithPadding(tokenizer=tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 优化器不需要创建了"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForSequenceClassification were not initialized from the model checkpoint at /data/models/huggingface/rbt3 and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "from torch.optim import AdamW\n",
    "model =  AutoModelForSequenceClassification.from_pretrained(pretrain_model_dir,num_labels=2)\n",
    "if torch.cuda.is_available():\n",
    "    model = model.cuda()\n",
    "\n",
    "optimizer =  AdamW(model.parameters(),lr=5e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 不需要创建优化器，只需要创建模型即可"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model =  AutoModelForSequenceClassification.from_pretrained('D:/code/models/huggingface/rbt3',num_labels=2)\n",
    "# cuda判定也不需要了\n",
    "# if torch.cuda.is_available():\n",
    "    # model = model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 评估部分重写evaluate的部分\n",
    "#### 1、在huggingface.co 网站中找tasks 然后在看支持的评估指标\n",
    "#### 2、直接下载失败，可以github上下载源代码，然后直接加载文件夹。\n",
    "#### 3、使用 add_batch 和 最终的compute方法最后实现评估指标的汇总工作\n",
    "#### 4、可以画出雷达图"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def evaluate():\n",
    "    model.eval()\n",
    "    acc_num = 0\n",
    "    with torch.inference_mode():\n",
    "        for batch in train_loader:\n",
    "            if torch.cuda.is_available():\n",
    "                batch = {k: v.cuda() for k, v in batch.items()}\n",
    "            output = model(**batch)\n",
    "            pred = torch.argmax(output.logits, dim=-1)\n",
    "            acc_num += (pred.long() == batch[\"labels\"].long()).float().sum()\n",
    "    return acc_num / len(validset)\n",
    "\n",
    "def train(epoch=3, log_step=100):\n",
    "    global_step = 0\n",
    "    for ep in range(epoch):\n",
    "        model.train()\n",
    "        start = time.time()\n",
    "        for batch in train_loader:\n",
    "            if torch.cuda.is_available():\n",
    "                batch = {k: v.cuda() for k, v in batch.items()}\n",
    "            optimizer.zero_grad()\n",
    "            output = model(**batch)\n",
    "            loss = output.loss.mean()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            if global_step % log_step == 0:\n",
    "                print(f\"ep: {ep}, global_step: {global_step}, loss: {loss.item()}\")\n",
    "            global_step += 1\n",
    "        acc = evaluate()\n",
    "        print(f\"ep: {ep}, acc: {acc}, time: {time.time() - start}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建 TrainingArguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ep: 0, global_step: 0, loss: 0.7963226437568665\n",
      "ep: 0, global_step: 100, loss: 0.27849486470222473\n",
      "ep: 0, global_step: 200, loss: 0.33037716150283813\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'validset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[33], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[1;32mIn[32], line 31\u001b[0m, in \u001b[0;36mtrain\u001b[1;34m(epoch, log_step)\u001b[0m\n\u001b[0;32m     29\u001b[0m         \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mep: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mep\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, global_step: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mglobal_step\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, loss: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mloss\u001b[38;5;241m.\u001b[39mitem()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     30\u001b[0m     global_step \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m---> 31\u001b[0m acc \u001b[38;5;241m=\u001b[39m \u001b[43mevaluate\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     32\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mep: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mep\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, acc: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00macc\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, time: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtime\u001b[38;5;241m.\u001b[39mtime()\u001b[38;5;250m \u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;250m \u001b[39mstart\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[1;32mIn[32], line 13\u001b[0m, in \u001b[0;36mevaluate\u001b[1;34m()\u001b[0m\n\u001b[0;32m     11\u001b[0m         pred \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39margmax(output\u001b[38;5;241m.\u001b[39mlogits, dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m     12\u001b[0m         acc_num \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m (pred\u001b[38;5;241m.\u001b[39mlong() \u001b[38;5;241m==\u001b[39m batch[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlabels\u001b[39m\u001b[38;5;124m\"\u001b[39m]\u001b[38;5;241m.\u001b[39mlong())\u001b[38;5;241m.\u001b[39mfloat()\u001b[38;5;241m.\u001b[39msum()\n\u001b[1;32m---> 13\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m acc_num \u001b[38;5;241m/\u001b[39m \u001b[38;5;28mlen\u001b[39m(\u001b[43mvalidset\u001b[49m)\n",
      "\u001b[1;31mNameError\u001b[0m: name 'validset' is not defined"
     ]
    }
   ],
   "source": [
    "train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 创建  Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate()\n",
    "trainer.evaluate(tokenizer_datasets['test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.predict(tokenizer_datasets['test'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 保存模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def save_model(model,tokenizer,epoch):\n",
    "#     original_path = \"D:/code/models/classfication_rbt3/\"\n",
    "#     path = os.path.join(original_path,'epoch_{}'.format(epoch))\n",
    "#     if not os.path.exists(path):\n",
    "#         os.makedirs(path)\n",
    "#     model.save_pretrained(path)\n",
    "#     tokenizer.save_pretrained(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def train(epoch=5,log_step=100):\n",
    "#     global_step = 0\n",
    "#     for epoch in range(epoch):\n",
    "#         print('epoch:{}'.format(epoch))\n",
    "#         model.train()\n",
    "#         for batch in train_loader:\n",
    "#             if torch.cuda.is_available():\n",
    "#                 batch = {k:v.cuda() for k,v in batch.items()}\n",
    "#             optimizer.zero_grad()\n",
    "#             output = model(**batch)\n",
    "#             output.loss.backward()\n",
    "#             optimizer.step()\n",
    "#             if global_step % log_step == 0:\n",
    "#                 print('epoch:{},step:{},loss:{}'.format(epoch,global_step,output.loss.item()))\n",
    "#             global_step+=1  \n",
    "#         acc = evaluate_model()\n",
    "#         print('epoch:{},acc:{}'.format(epoch,acc))\n",
    "#         save_model(model,tokenizer,epoch)\n",
    "#     return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# acc = train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(type(acc),acc.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 评估结果的可视化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from evaluate.visualization import radar_plot\n",
    "\n",
    "test_out =  trainer.evaluate(tokenizer_datasets['test'])\n",
    "print(test_out)\n",
    "\n",
    "radar_out = [{\n",
    "        \"accuracy\": test_out['eval_accuracy'],\n",
    "        \"f1\":  test_out['eval_f1'],\n",
    "        \"precision\":  test_out['eval_precision'],\n",
    "        \"recall\":  test_out['eval_recall']\n",
    "}]\n",
    "\n",
    "model_names = [\n",
    "    \"model 1\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot =  radar_plot(data=radar_out,model_names=model_names)\n",
    "plot.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
